{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9bcfbc65",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/yuhangya/Library/Python/3.8/lib/python/site-packages/cvxpy/problems/problem.py:1278: UserWarning: Solution may be inaccurate. Try another solver, adjusting the solver settings, or solve with verbose=True for more information.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "ename": "SolverError",
     "evalue": "Solver 'ECOS' failed. Try another solver, or solve with verbose=True for more information.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mSolverError\u001b[0m                               Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-3-4032c2378609>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     77\u001b[0m         \u001b[0mrho_ij\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msqrt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m3\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlog\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mt\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;36m2\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mT_ij\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mrs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     78\u001b[0m         \u001b[0mmu_bar\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmaximum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmu_hat\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mrho_ij\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# LCB\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 79\u001b[0;31m         \u001b[0mx_t\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstatus\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0moracle\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmu_bar\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     80\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0;34m'optimal'\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mstatus\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     81\u001b[0m             \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'Solution infeasible 2'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-3-4032c2378609>\u001b[0m in \u001b[0;36moracle\u001b[0;34m(y, mu)\u001b[0m\n\u001b[1;32m     33\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     34\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0moracle\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmu\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 35\u001b[0;31m     \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprob\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0moptimization\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mm\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwS\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmu\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mBS\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbeta\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     36\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprob\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprob\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstatus\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     37\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-3-4032c2378609>\u001b[0m in \u001b[0;36moptimization\u001b[0;34m(m, n, wS, y, mu, BS, beta)\u001b[0m\n\u001b[1;32m     29\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     30\u001b[0m     \u001b[0mprob\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mProblem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mMinimize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconstraints\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 31\u001b[0;31m     \u001b[0mprob\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msolve\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# Returns the optimal value.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     32\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprob\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     33\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Library/Python/3.8/lib/python/site-packages/cvxpy/problems/problem.py\u001b[0m in \u001b[0;36msolve\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    460\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    461\u001b[0m             \u001b[0msolve_func\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mProblem\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_solve\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 462\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0msolve_func\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    463\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    464\u001b[0m     \u001b[0;34m@\u001b[0m\u001b[0mclassmethod\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Library/Python/3.8/lib/python/site-packages/cvxpy/problems/problem.py\u001b[0m in \u001b[0;36m_solve\u001b[0;34m(self, solver, warm_start, verbose, gp, qcp, requires_grad, enforce_dpp, **kwargs)\u001b[0m\n\u001b[1;32m    960\u001b[0m         \u001b[0mend\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtime\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    961\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_solve_time\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mend\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mstart\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 962\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munpack_results\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msolution\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msolving_chain\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minverse_data\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    963\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mverbose\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    964\u001b[0m             \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0m_FOOTER\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Library/Python/3.8/lib/python/site-packages/cvxpy/problems/problem.py\u001b[0m in \u001b[0;36munpack_results\u001b[0;34m(self, solution, chain, inverse_data)\u001b[0m\n\u001b[1;32m   1282\u001b[0m             )\n\u001b[1;32m   1283\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0msolution\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstatus\u001b[0m \u001b[0;32min\u001b[0m \u001b[0ms\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mERROR\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1284\u001b[0;31m             raise error.SolverError(\n\u001b[0m\u001b[1;32m   1285\u001b[0m                     \u001b[0;34m\"Solver '%s' failed. \"\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0mchain\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msolver\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1286\u001b[0m                     \u001b[0;34m\"Try another solver, or solve with verbose=True for more \"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mSolverError\u001b[0m: Solver 'ECOS' failed. Try another solver, or solve with verbose=True for more information."
     ]
    }
   ],
   "source": [
    "# soft constraints\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import cvxpy as cp\n",
    "\n",
    "N = 1\n",
    "T = 500\n",
    "m = 20  # number of devices\n",
    "n = 5  # number of servers\n",
    "y_max = 110\n",
    "y_min = 90\n",
    "rs = 0.1\n",
    "beta = np.ones(n)\n",
    "beta[0] = 3\n",
    "beta[1] = 3\n",
    "# beta[3] = 5\n",
    "gamma = 0.0001\n",
    "seed = 0\n",
    "np.random.seed(seed)\n",
    "\n",
    "def optimization(m, n, wS, y, mu, BS, beta):\n",
    "    x = cp.Variable((n + 1, m), nonneg=True)\n",
    "    obj = beta @ cp.inv_pos(cp.sqrt(x[1:, :] @ y + wS)) + gamma * cp.sum(cp.multiply(cp.diag(mu @ x[1:, :]), y))\n",
    "\n",
    "    constraints = [0 <= x, x <= 1,\n",
    "                   x[1:, :] @ y <= BS,\n",
    "                   cp.sum(x, 0) == 1]\n",
    "\n",
    "    prob = cp.Problem(cp.Minimize(obj), constraints)\n",
    "    prob.solve()  # Returns the optimal value.\n",
    "    return x, prob\n",
    "\n",
    "def oracle(y, mu):\n",
    "    x, prob = optimization(m, n, wS, y, mu, BS, beta)\n",
    "    return x.value, prob.value, prob.status\n",
    "\n",
    "def f(x, y, mu):\n",
    "    return beta.dot(1/np.sqrt(x[1:, :].dot(y) + wS)) + gamma * np.sum((y*x[1:]).T*mu)\n",
    "\n",
    "def f_drop(x, y, mu, BS, gamma):\n",
    "    dp = np.minimum(x[1:, :].dot(y), BS)\n",
    "    return beta.dot(1/np.sqrt(dp + wS)) + gamma * np.sum((y*x[1:]).T*mu)\n",
    "\n",
    "reg = np.zeros((N, T))\n",
    "\n",
    "#records all y, x_opt, x_t #yuhang yao\n",
    "y_N_T = np.zeros((N, T, m)) #yuhang yao\n",
    "x_opt_N_T = np.zeros((N, T, n + 1, m)) #yuhang yao\n",
    "x_t_N_T = np.zeros((N, T, n + 1, m)) #yuhang yao\n",
    "j_N_T = np.zeros((N, T, m))#yuhang yao\n",
    "BS_N = np.zeros((N, n))#yuhang yao\n",
    "\n",
    "for u in range(N):\n",
    "    wS = np.random.randint(15, 25, n)\n",
    "    BS = np.random.uniform(y_min*10, y_max*10, n)\n",
    "    BS_N[u] = BS\n",
    "    mu = np.random.rand(m, n)\n",
    "    mu[:,0] = 0.5\n",
    "    mu[:,1] = 0.5\n",
    "    mu[:,2] = 0.8\n",
    "    mu[:,3] = 0.8\n",
    "    mu[:,4] = 0.8\n",
    "    # trace_gen = Trace(m, n, seed + u)\n",
    "    # mu = trace_gen.avg()\n",
    "    # mu = np.random.rand(m, n)\n",
    "    # mu_hat = np.zeros_like(mu)  # empirical mean\n",
    "    mu_hat = np.zeros_like(mu)\n",
    "    T_ij = np.ones_like(mu)  # total number of times arm (i,j) is played\n",
    "    for t in range(T):\n",
    "        y = np.random.uniform(y_min, y_max, m).astype(int)\n",
    "        x_opt, f_opt, status = oracle(y, mu)\n",
    "        if 'optimal' not in status:\n",
    "            print('Solution infeasible 1')\n",
    "            break\n",
    "\n",
    "        rho_ij = np.sqrt(3 * np.log(t + 1) / (2 * T_ij)) * rs\n",
    "        mu_bar = np.maximum(mu_hat - rho_ij, 0) # LCB\n",
    "        x_t, _, status = oracle(y, mu_bar)\n",
    "        if 'optimal' not in status:\n",
    "            print('Solution infeasible 2')\n",
    "            break\n",
    "\n",
    "        f_t = f(x_t, y, mu)\n",
    "\n",
    "        # sample j based on x_t[i], observe c_ij, update mu_hat[i,j]\n",
    "        # c = trace_gen.generate()\n",
    "        for i in range(m):\n",
    "            j = np.random.choice(n+1, p=x_t[:, i])\n",
    "            j_N_T[u, t, i] = j #yuhang yao\n",
    "            if j != 0:\n",
    "                j -= 1\n",
    "                c_ij = int(np.random.rand() < mu[i, j])\n",
    "                # a = np.random.rand() * 3\n",
    "                # c_ij = np.random.beta(a, a * (1-mu[i, j])/mu[i, j]) # beta distribution\n",
    "                # c_ij = c[i, j]  # trace\n",
    "                T_ij[i, j] += 1\n",
    "                mu_hat[i, j] += (c_ij - mu_hat[i, j]) / T_ij[i, j]\n",
    "\n",
    "        # calculate regert\n",
    "        reg[u, t] = f_t - f_opt\n",
    "        y_N_T[u, t] = y#yuhang yao\n",
    "        x_opt_N_T[u, t] = x_opt#yuhang yao\n",
    "        x_t_N_T[u, t] = x_t#yuhang yao\n",
    "        \n",
    "plt.plot(np.cumsum(reg, axis=1).T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "chemical-specification",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "chemical-luther",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "np.save(\"soft_wS\", wS)\n",
    "np.save(\"soft_BS_N\", BS_N)\n",
    "np.save(\"soft_y_N_T\", y_N_T)\n",
    "np.save(\"soft_x_opt_N_T\", x_opt_N_T)\n",
    "np.save(\"soft_x_t_N_T\", x_t_N_T)\n",
    "np.save(\"soft_j_N_T\", j_N_T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "applicable-watch",
   "metadata": {},
   "outputs": [],
   "source": [
    "from functions import *\n",
    "wS = np.load(\"soft_wS.npy\")\n",
    "y_N_T = np.load(\"soft_y_N_T.npy\")\n",
    "x_opt_N_T = np.load(\"soft_x_opt_N_T.npy\")\n",
    "x_t_N_T = np.load(\"soft_x_t_N_T.npy\")\n",
    "j_N_T = np.load(\"soft_j_N_T.npy\")\n",
    "BS_N = np.load(\"soft_BS_N.npy\")\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fifth-indication",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "wS = wS.astype(int)\n",
    "BS_N = BS_N.astype(int)\n",
    "y_N_T = y_N_T.astype(int)\n",
    "\n",
    "j_N_T = j_N_T.astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2061cedb",
   "metadata": {},
   "outputs": [],
   "source": [
    "Data_num_S_N_T = np.zeros((N, T, wS.shape[0]))\n",
    "for u in range(N):\n",
    "    for t in range(T):\n",
    "        Data_num_S_N_T[u, t] = wS\n",
    "        for i in range(m):\n",
    "            \n",
    "                    j = j_N_T[u, t, i]#np.random.choice(n + 1, p=x_t_N_T[u, t, :, i])\n",
    "                    \n",
    "                    if j != 0:\n",
    "                        j -= 1\n",
    "                        Data_num_S_N_T[u, t, j] += y_N_T[u, t, i] #upload to device j-1\n",
    "                        Data_num_S_N_T[u, t, j] = min(Data_num_S_N_T[u, t, j], BS_N[u, j])\n",
    "\n",
    "\n",
    "from fed_train import fed_T\n",
    "import time\n",
    "save_time = str(time.time())\n",
    "import torch\n",
    "torch.manual_seed(seed)\n",
    "np.random.seed(seed)\n",
    "\n",
    "total_loss_train = []\n",
    "total_accuracy_train = []\n",
    "\n",
    "total_acc_test = []\n",
    "total_loss_test = []\n",
    "\n",
    "for u in range(N):\n",
    "    L_T = []\n",
    "    for t in range(T):\n",
    "        L = list(Data_num_S_N_T[u, t])\n",
    "        L_T.append(L)\n",
    "    L_T = np.array(L_T).astype(int)\n",
    "\n",
    "\n",
    "\n",
    "    loss_train, accuracy_train, acc_tests, loss_tests, args = fed_T(L_T, \"mnist\", False)\n",
    "    if len(total_loss_train) == 0:\n",
    "        total_loss_train = np.array(loss_train)\n",
    "    else:\n",
    "        total_loss_train += np.array(loss_train)\n",
    "        \n",
    "    if len(total_accuracy_train) == 0:\n",
    "        total_accuracy_train = np.array(accuracy_train)\n",
    "    else:\n",
    "        total_accuracy_train += np.array(accuracy_train)\n",
    "        \n",
    "        \n",
    "    if len(total_acc_test) == 0:\n",
    "        total_acc_test = np.array(acc_tests)\n",
    "    else:\n",
    "        total_acc_test += np.array(acc_tests)\n",
    "    \n",
    "    if len(total_loss_test) == 0:\n",
    "        total_loss_test = np.array(loss_tests)\n",
    "    else:\n",
    "        total_loss_test += np.array(loss_tests)\n",
    "\n",
    "    \n",
    "total_loss_train /= N\n",
    "total_accuracy_train /= N\n",
    "total_acc_test /= N\n",
    "total_loss_test /= N\n",
    "\n",
    "np.save(\"soft_iid\" + str(args.iid) + \"total_loss_train\", total_loss_train)\n",
    "np.save(\"soft_iid\" + str(args.iid) + \"total_accuracy_train\", total_accuracy_train)\n",
    "np.save(\"soft_iid\" + str(args.iid) + \"total_acc_test\", total_acc_test)\n",
    "np.save(\"soft_iid\" + str(args.iid) + \"total_loss_test\", total_loss_test)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "protected-state",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(seed)\n",
    "np.random.seed(seed)\n",
    "Data_num_S_N_T = np.zeros((N, T, wS.shape[0]))\n",
    "for u in range(N):\n",
    "    for t in range(T):\n",
    "        Data_num_S_N_T[u, t] = wS\n",
    "        for i in range(m):\n",
    "            \n",
    "                    j = np.random.choice(n + 1)\n",
    "                    \n",
    "                    if j != 0:\n",
    "                        j -= 1\n",
    "                        Data_num_S_N_T[u, t, j] += y_N_T[u, t, i] #upload to device j-1\n",
    "                        Data_num_S_N_T[u, t, j] = min(Data_num_S_N_T[u, t, j], BS_N[u, j])\n",
    "\n",
    "\n",
    "from fed_train import fed_T\n",
    "import time\n",
    "save_time = str(time.time())\n",
    "import torch\n",
    "\n",
    "\n",
    "total_loss_train = []\n",
    "total_accuracy_train = []\n",
    "\n",
    "total_acc_test = []\n",
    "total_loss_test = []\n",
    "\n",
    "for u in range(N):\n",
    "    L_T = []\n",
    "    for t in range(T):\n",
    "        L = list(Data_num_S_N_T[u, t])\n",
    "        L_T.append(L)\n",
    "    L_T = np.array(L_T).astype(int)\n",
    "\n",
    "\n",
    "\n",
    "    loss_train, accuracy_train, acc_tests, loss_tests, args = fed_T(L_T, \"mnist\", False)\n",
    "    if len(total_loss_train) == 0:\n",
    "        total_loss_train = np.array(loss_train)\n",
    "    else:\n",
    "        total_loss_train += np.array(loss_train)\n",
    "        \n",
    "    if len(total_accuracy_train) == 0:\n",
    "        total_accuracy_train = np.array(accuracy_train)\n",
    "    else:\n",
    "        total_accuracy_train += np.array(accuracy_train)\n",
    "        \n",
    "        \n",
    "    if len(total_acc_test) == 0:\n",
    "        total_acc_test = np.array(acc_tests)\n",
    "    else:\n",
    "        total_acc_test += np.array(acc_tests)\n",
    "    \n",
    "    if len(total_loss_test) == 0:\n",
    "        total_loss_test = np.array(loss_tests)\n",
    "    else:\n",
    "        total_loss_test += np.array(loss_tests)\n",
    "\n",
    "    \n",
    "total_loss_train /= N\n",
    "total_accuracy_train /= N\n",
    "total_acc_test /= N\n",
    "total_loss_test /= N\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "np.save(\"random_soft_iid\" + str(args.iid) + \"total_loss_train\", total_loss_train)\n",
    "np.save(\"random_soft_iid\" + str(args.iid) + \"total_accuracy_train\", total_accuracy_train)\n",
    "np.save(\"random_soft_iid\" + str(args.iid) + \"total_acc_test\", total_acc_test)\n",
    "np.save(\"random_soft_iid\" + str(args.iid) + \"total_loss_test\", total_loss_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "civil-enough",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
